import math

from pandas import option_context
from sympy import true
from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np
import math
from .padding_resize_convert import PaddingResizeConvert

#切补丁
class SplicingConvert(ConvertBase):
    # 初始化
    #图片数量 img_count : 必须是n²的数量
    def __init__(self, img_count = 4, use_rate=0.5, is_resize = False, filter_small_box_rate = 0.25):
        super().__init__(use_rate=use_rate)
        #图片数量
        self.img_count = img_count
        self.horization_count = math.sqrt(img_count)
        self.vertical_count = math.sqrt(img_count)
        
        self.filter_small_box_rate = filter_small_box_rate # 是否过滤较小的box 过滤的阈值
        self.is_resize = is_resize
        self.padding_resize = PaddingResizeConvert((512, 512))
    
    def __get_range(self, img_size, range_size):
        return int(np.random.randint(0, img_size[0] - range_size[0])), int(np.random.randint(0, img_size[1] - range_size[1]))
    
    def __convert_boxes(self, org_boxes, splices, start_point):
        boxes = org_boxes.copy()
        boxes[:, [0, 2]] = np.minimum(boxes[:, [0, 2]], splices[2])
        boxes[:, [0, 2]] = np.maximum(boxes[:, [0, 2]] - splices[0], 0)
        
        boxes[:, [1, 3]] = np.minimum(boxes[:, [1, 3]], splices[3])
        boxes[:, [1, 3]] = np.maximum(boxes[:, [1, 3]] - splices[1], 0)
        
        res_boxes = []
        # for box in boxes:
        #     if box[0] == box[2] or box[1] == box[3]:
        #         continue
            
        #     box[[0, 2]] = box[[0, 2]] + start_point[0]
        #     box[[1, 3]] = box[[1, 3]] + start_point[1]
        #     res_boxes.append(box)
        for org, cur in zip(*(org_boxes, boxes)):
            area_org = (org[2] - org[0]) * (org[3] - org[1])
            area_cur = (cur[2] - cur[0]) * (cur[3] - cur[1])

            if area_cur / area_org <= self.filter_small_box_rate:
                continue
            
            cur[[0, 2]] = cur[[0, 2]] + start_point[0]
            cur[[1, 3]] = cur[[1, 3]] + start_point[1]
            res_boxes.append(cur)
        
        return res_boxes
        
    
    def __convert_points(self, points):
        pass
    
    def convert(self, imgs, boxes, points, ranges = None):
        assert len(imgs) == self.img_count
        self.horization_count = math.sqrt(self.img_count)
        self.vertical_count = math.sqrt(self.img_count)
        
        if self.is_resize == False:
            return self.__clip_convert(imgs, boxes, points, ranges)
        else:
            return self.__cat_convert(imgs, boxes, points)
    
    def __cat_convert(self, imgs, boxes, points):
        shape = imgs[0].shape
        h, w, _ = shape
        #补丁的宽和高
        patches_w = int(w // self.horization_count)
        patches_h = int(h // self.vertical_count)
        target_img = np.zeros((h, w, 3), dtype=np.uint8)
        # 设置缩放尺寸
        self.padding_resize.targetsize = (patches_w, patches_h)
        #当前拼接补丁的左上角的位置
        start_x = 0
        start_y = 0
        res_boxes = []
        
        for i in range(len(imgs)):
            img = imgs[i]
            if boxes is not None:
                box = boxes[i]
                if box.shape[1] > 4:
                    cs = box[:, -1]
                else:
                    cs = None
            else:
                box = None
            if points is not None:
                point = points[i]
            else:
                point = None
            
            img, box, point = self.padding_resize.convert(img, box, point)
            
            target_img[start_y: start_y + patches_h, start_x: start_x + patches_w] = img
            if box is not None:
                box[:, [0, 2]] += start_x
                box[:, [1, 3]] += start_y
                if cs is not None:
                    box[:, -1] = cs
                for b in box:
                    res_boxes.append(b)
            if point is not None:
                point += np.array([start_x, start_y], np.float32)
                points[i] = point            
            
            # 更新当前拼接补丁的左上角的位置
            start_x += patches_w
            if start_x >= w or (w - start_x) < patches_w:
                start_x = 0
                start_y += patches_h
                
                if (h - start_y) < patches_h:
                    break
        
        return target_img, np.array(res_boxes, np.float32), None, None
    
        
    def __clip_convert(self, imgs, boxes, points, ranges = None):
        shape = imgs[0].shape
        h, w, _ = shape
        # 操作区域集合
        option_ranges = []
        #补丁的宽和高
        patches_w = int(w // self.horization_count)
        patches_h = int(h // self.vertical_count)
        target_img = np.zeros((h, w, 3), dtype=np.uint8)
        #当前拼接补丁的左上角的位置
        start_x = 0
        start_y = 0
        
        res_boxes = []
        
        for i in range(len(imgs)):
            img = imgs[i]
            if ranges is None:
                range_start_x, range_start_y = self.__get_range((w, h), (patches_w, patches_h))
            else:
                range_start_x, range_start_y = ranges[i]
            #记录操作区域
            option_ranges.append([range_start_x, range_start_y])
            # print(range_start_x, range_start_y, start_x, start_y, patches_w, patches_h)
            target_img[start_y: start_y + patches_h, start_x: start_x + patches_w] = img[range_start_y: range_start_y + patches_h, range_start_x: range_start_x + patches_w]
            if boxes is not None:
                #转换对应的 坐标和点
                img_boxes = boxes[i]
                rg = [range_start_x, range_start_y, range_start_x + patches_w, range_start_y + patches_h]
                res = self.__convert_boxes(img_boxes, rg, (start_x, start_y))
                if len(res) > 0:
                    for b in res:
                        res_boxes.append(b)
                        
                #  不支持多点转换
                # points[i] = self.__convert_points(img_points)
            
            # 更新当前拼接补丁的左上角的位置
            start_x += patches_w
            if start_x >= w:
                start_x = 0
                start_y += patches_h
        
        return target_img, np.array(res_boxes, np.float32), None, option_ranges
    
    